"""
LCEL Routing Chain Service

This service demonstrates LangChain Expression Language (LCEL) for routing chains.
LCEL is the modern, recommended way to build chains in LangChain.
"""

from typing import Dict, List, Any, Literal, Union
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnableBranch, RunnableLambda, RunnablePassthrough
from langchain_openai import ChatOpenAI
from langchain.schema.runnable import RunnableConfig
from langchain_core.messages import HumanMessage, AIMessage

from app.services.zodiac_service import get_zodiac_sign
from app.services.system_agent import handle_system_query
from app.core.config import settings


class LCELRoutingService:
    """
    LCEL Routing Service using modern LangChain Expression Language
    """

    def __init__(self, model_name: str = "deepseek-chat", temperature: float = 0.7):
        self.model_name = model_name
        self.temperature = temperature
        self.llm = ChatOpenAI(
            model=model_name,
            temperature=temperature,
            max_tokens=65536,
            api_key=settings.openai_api_key,
            base_url=settings.openai_base_url
        )
        self.parser = StrOutputParser()

    def create_zodiac_chain_lcel(self):
        """Create zodiac chain using LCEL syntax"""
        # Create zodiac processing function using RunnableLambda
        zodiac_processor = RunnableLambda(
            lambda x: self._process_zodiac_query(x["query"] if isinstance(x, dict) else x)
        )

        # LCEL chain: processor | parser
        zodiac_chain = zodiac_processor | self.parser

        return zodiac_chain

    def create_system_chain_lcel(self):
        """Create system chain using LCEL syntax"""
        # Create system processing function using RunnableLambda
        system_processor = RunnableLambda(
            lambda x: self._process_system_query(x["query"] if isinstance(x, dict) else x)
        )

        # LCEL chain: processor | parser
        system_chain = system_processor | self.parser

        return system_chain

    def create_calculation_chain_lcel(self):
        """Create calculation chain using LCEL syntax"""
        # Create calculation processing function using RunnableLambda
        calc_processor = RunnableLambda(
            lambda x: self._process_calculation_query(x["query"] if isinstance(x, dict) else x)
        )

        # LCEL chain: processor | parser
        calc_chain = calc_processor | self.parser

        return calc_chain

    def create_general_chain_lcel(self):
        """Create general conversation chain using LCEL syntax"""
        general_prompt = ChatPromptTemplate.from_messages([
            ("system", """你是一个智能助手，可以回答各种问题。
你能够：
- 进行日常对话
- 回答常识性问题
- 提供有用的信息和建议
- 解释概念和定义

请用友好、自然的方式与用户交流。记住对话的上下文，提供连贯的回答。"""),
            ("human", "{query}")
        ])

        # Simple LCEL chain: prompt | llm | parser
        general_chain = general_prompt | self.llm | self.parser

        return general_chain

    def create_memory_aware_general_chain_lcel(self):
        """Create memory-aware general conversation chain using LCEL syntax"""
        general_prompt = ChatPromptTemplate.from_messages([
            ("system", """你是一个智能助手，可以回答各种问题。
你能够：
- 进行日常对话
- 回答常识性问题
- 提供有用的信息和建议
- 解释概念和定义

请用友好、自然的方式与用户交流。重要的是要记住对话的上下文历史，
提供连贯、有上下文的回答。如果用户提到之前的内容，请参考历史记录。"""),
            MessagesPlaceholder(variable_name="conversation_history"),
            ("human", "{query}")
        ])

        # LCEL chain with memory: prompt | llm | parser
        general_chain = general_prompt | self.llm | self.parser

        return general_chain

    def create_routing_branch_lcel(self):
        """
        Create routing using RunnableBranch - the LCEL way
        """
        # Define individual chains
        zodiac_chain = self.create_zodiac_chain_lcel()
        system_chain = self.create_system_chain_lcel()
        calc_chain = self.create_calculation_chain_lcel()
        general_chain = self.create_general_chain_lcel()

        # Use RunnableBranch for routing (LCEL style)
        routing_branch = RunnableBranch(
            (lambda x: self._is_zodiac_query(x), zodiac_chain),
            (lambda x: self._is_system_query(x), system_chain),
            (lambda x: self._is_calculation_query(x), calc_chain),
            general_chain  # Default branch
        )

        return routing_branch

    def create_memory_aware_routing_branch_lcel(self):
        """
        Create memory-aware routing using RunnableBranch that supports conversation history
        """
        # Define individual chains with memory support
        zodiac_chain = self.create_zodiac_chain_lcel()
        system_chain = self.create_system_chain_lcel()
        calc_chain = self.create_calculation_chain_lcel()
        general_chain = self.create_memory_aware_general_chain_lcel()

        # Use RunnableBranch for routing (LCEL style)
        routing_branch = RunnableBranch(
            (lambda x: self._is_zodiac_query(x["query"]), zodiac_chain),
            (lambda x: self._is_system_query(x["query"]), system_chain),
            (lambda x: self._is_calculation_query(x["query"]), calc_chain),
            general_chain  # Default branch (memory-aware)
        )

        return routing_branch

    def create_composed_routing_lcel(self):
        """
        Create composed routing using LCEL composition
        """
        # Define the chains
        chains = {
            "zodiac": self.create_zodiac_chain_lcel(),
            "system": self.create_system_chain_lcel(),
            "calculation": self.create_calculation_chain_lcel(),
            "general": self.create_general_chain_lcel()
        }

        # Create routing logic
        def route_and_execute(query: str) -> str:
            # Determine route
            route = self._get_route_type(query)

            # Execute corresponding chain
            selected_chain = chains[route]
            return selected_chain.invoke(query)

        # Wrap in RunnableLambda for LCEL compatibility
        routing_composed = RunnableLambda(route_and_execute)

        return routing_composed

    def create_parallel_routing_lcel(self):
        """
        Create parallel routing using LCEL's parallel capabilities
        """
        # Define specialized chains
        zodiac_chain = self.create_zodiac_chain_lcel()
        system_chain = self.create_system_chain_lcel()
        calc_chain = self.create_calculation_chain_lcel()

        # Create classification chain
        classifier_prompt = ChatPromptTemplate.from_messages([
            ("system", """你是一个查询分类专家。请分析用户的查询，将其分类为以下类型之一：
- zodiac: 生肖相关查询
- system: 系统监控相关查询
- calculation: 数学计算相关查询
- general: 一般性对话

只返回分类名称，不要其他内容。"""),
            ("human", "查询: {query}\n分类:")
        ])

        classifier_chain = classifier_prompt | self.llm | self.parser

        # Create parallel execution
        parallel_branch = RunnableBranch(
            (lambda x: x["classification"] == "zodiac", zodiac_chain),
            (lambda x: x["classification"] == "system", system_chain),
            (lambda x: x["classification"] == "calculation", calc_chain),
            self.create_general_chain_lcel()
        )

        # Compose the complete chain
        parallel_routing_chain = (
            {
                "query": RunnablePassthrough(),
                "classification": classifier_chain
            }
            | parallel_branch
        )

        return parallel_routing_chain

    # Helper methods for routing decisions
    def _decide_route_lcel(self, query: str) -> str:
        """Route decision function for LCEL"""
        return self._get_route_type(query)

    def _get_route_type(self, query: str) -> str:
        """Get route type based on query content"""
        query_lower = query.lower()

        zodiac_keywords = ["生肖", "属相", "鼠", "牛", "虎", "兔", "龙", "蛇", "马", "羊", "猴", "鸡", "狗", "猪", "本命年", "配对"]
        system_keywords = ["系统", "服务器", "负载", "会话", "统计", "数据库", "性能", "内存", "cpu", "磁盘", "监控"]
        calculation_keywords = ["计算", "等于", "加", "减", "乘", "除", "百分比", "多少", "数学"]

        if any(keyword in query_lower for keyword in zodiac_keywords):
            return "zodiac"
        elif any(keyword in query_lower for keyword in system_keywords):
            return "system"
        elif any(keyword in query_lower for keyword in calculation_keywords):
            return "calculation"
        else:
            return "general"

    def _is_zodiac_query(self, query: str) -> bool:
        return self._get_route_type(query) == "zodiac"

    def _is_system_query(self, query: str) -> bool:
        return self._get_route_type(query) == "system"

    def _is_calculation_query(self, query: str) -> bool:
        return self._get_route_type(query) == "calculation"

    def _format_conversation_history(self, conversation_history: List[Dict[str, Any]]) -> List[Union[HumanMessage, AIMessage]]:
        """
        Format conversation history for LangChain message format
        """
        formatted_history = []

        for msg in conversation_history:
            if msg.get("role") == "user":
                formatted_history.append(HumanMessage(content=msg.get("content", "")))
            elif msg.get("role") == "assistant":
                formatted_history.append(AIMessage(content=msg.get("content", "")))

        return formatted_history

    # Processing methods for different query types
    def _process_zodiac_query(self, query: str) -> str:
        """Process zodiac query with integrated function calls"""
        import re

        # Try to extract year from query
        year_match = re.search(r'\b(19|20)\d{2}\b', query)
        if year_match:
            year = int(year_match.group())
            zodiac_info = get_zodiac_sign(year)
            if "error" not in zodiac_info:
                return f"""
根据您的查询，以下是生肖信息：

🐉 生肖: {zodiac_info['zodiac_sign']} ({zodiac_info['zodiac_english']})
🐲 生肖年: {zodiac_info['zodiac_year']}年
⭐ 特质: {zodiac_info['attributes']['traits']}
🎰 幸运数字: {', '.join(map(str, zodiac_info['attributes']['lucky_numbers']))}
🎨 幸运颜色: {', '.join(zodiac_info['attributes']['lucky_colors'])}

如果您需要更详细的信息或其他年份的生肖，请告诉我！
                """.strip()

        # If no year found or general zodiac question, use LLM
        return f"请回答以下关于生肖的问题: {query}"

    def _process_system_query(self, query: str) -> str:
        """Process system query with integrated function calls"""
        try:
            # Use system agent to get real system information
            system_info = handle_system_query(query)
            return system_info
        except Exception as e:
            return f"获取系统信息时出现错误，请告诉我您想了解的具体系统信息: {query}"

    def _process_calculation_query(self, query: str) -> str:
        """Process calculation query with integrated function calls"""
        import re

        # Try to evaluate simple expressions
        expr_match = re.search(r'([\d+\-*/().\s]+)', query)
        if expr_match:
            expr = expr_match.group(1).strip()
            try:
                # Safe evaluation
                result = eval(expr, {'__builtins__': None}, {})
                return f"计算结果: {expr} = {result}\n\n如果您需要更复杂的数学计算或解释，请告诉我！"
            except:
                pass

        # Fallback to LLM for complex calculations
        return f"请解答以下数学问题: {query}"

    def demo_all_lcel_approaches(self, query: str) -> Dict[str, str]:
        """
        Demo all LCEL routing approaches
        """
        results = {}

        # Test 1: RunnableBranch approach
        try:
            branch_chain = self.create_routing_branch_lcel()
            results["runnable_branch"] = branch_chain.invoke(query)
        except Exception as e:
            results["runnable_branch"] = f"Error: {str(e)}"

        # Test 2: Composed routing approach
        try:
            composed_chain = self.create_composed_routing_lcel()
            results["composed_routing"] = composed_chain.invoke(query)
        except Exception as e:
            results["composed_routing"] = f"Error: {str(e)}"

        # Test 3: Parallel routing approach
        try:
            parallel_chain = self.create_parallel_routing_lcel()
            results["parallel_routing"] = parallel_chain.invoke(query)
        except Exception as e:
            results["parallel_routing"] = f"Error: {str(e)}"

        return results

    def performance_test(self, query: str, iterations: int = 10) -> Dict[str, float]:
        """
        Performance test for different LCEL approaches
        """
        import time

        chains = {
            "runnable_branch": self.create_routing_branch_lcel(),
            "composed_routing": self.create_composed_routing_lcel(),
            "parallel_routing": self.create_parallel_routing_lcel()
        }

        results = {}

        for name, chain in chains.items():
            start_time = time.time()
            for _ in range(iterations):
                try:
                    chain.invoke(query)
                except Exception:
                    pass
            end_time = time.time()
            results[name] = end_time - start_time

        return results


# Global service instance
lcel_routing_service = LCELRoutingService()


def demo_lcel_approaches(query: str) -> Dict[str, str]:
    """
    Convenience function to demo LCEL approaches
    """
    return lcel_routing_service.demo_all_lcel_approaches(query)